import json
import logging
import os
from typing import List

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings
from langchain_postgres import PGVector
from langchain_text_splitters import CharacterTextSplitter

from app.config.settings import Setting

logger = logging.getLogger(__name__)


class Retriever:
    """基础检索器类，提供向量数据库检索功能。"""

    def __init__(self, collection_name: str = "chatbot"):
        self.embeddings = OllamaEmbeddings(base_url="http://localhost:11434", model="dengcao/Qwen3-Embedding-0.6B:Q8_0")

        self.embedding_filter = EmbeddingsFilter(embeddings=self.embeddings, similarity_threshold=0.5)

        self.pg_vector = PGVector(
            connection=PGVector.connection_string_from_db_params(
                driver="psycopg",
                host=Setting.pg_host,
                port=int(Setting.pg_port),
                user=Setting.pg_user,
                password=Setting.pg_password,
                database=Setting.pg_database,
            ),
            embeddings=self.embeddings,
            collection_name=collection_name,
        )
        self.retriever = self.pg_vector.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={
                'score_threshold': 0.1,
                'k': 5
            }
        )
        self.compression_retriever = ContextualCompressionRetriever(
            base_compressor=self.embedding_filter, base_retriever=self.retriever
        )
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=2000,
            chunk_overlap=200,
            separators=["\n\n", "\n", "。"],
        )
        self.docs_resolver = docs_resolve

    def load_data_from_txt(self, data_dir: str = "./data", suffix: str = ".txt"):
        """
        从文本文件中初始化数据。
        假设文本文件位于当前目录下的 "data" 文件夹中。
        """
        if not os.path.exists(data_dir):
            logger.error(f"数据目录 '{data_dir}' 不存在")
            return
        documents = []

        ##### TODO 为了避免重复加载，这里我们先删除现的向量集合，生产环境根据实际情况调整
        self.pg_vector.delete_collection()
        self.pg_vector.create_collection()
        ####################################

        for filename in os.listdir(data_dir):
            if filename.endswith(suffix):
                file_path = os.path.join(data_dir, filename)
                documents.extend(self.split_documents(file_path))
        if documents:
            # 自定义文档解析器
            if self.docs_resolver:
                documents = self.docs_resolver(documents)
                # 添加到向量数据库
            self.retriever.add_documents(documents)
            logger.info(f"成功添加 {len(documents)} 个文档到向量数据库")
        else:
            logger.warning("没有找到任何文档")

    def load_data_from_json(self, data_dir: str = "./data", suffix: str = ".json"):
        """
        从 JSON 文件中初始化数据。
        假设 JSON 文件位于当前目录下的 "data" 文件夹中。
        每个 JSON 文件包含一个列表，其中每个元素是一个字典，格式如下：
        [
            {
                "question": "查询航班动态流程",
                "answer": "1. 访问航空公司官网或APP，输入航班号查询动态。\n..."
            },
            ...
        ]
        """
        if not os.path.exists(data_dir):
            logger.error(f"数据目录 '{data_dir}' 不存在")
            return

        documents = []

        # 清空现有的向量集合（生产环境请根据实际情况调整）
        self.pg_vector.delete_collection()
        self.pg_vector.create_collection()

        for filename in os.listdir(data_dir):
            if filename.endswith(suffix):
                file_path = os.path.join(data_dir, filename)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)  # 加载 JSON 数据
                        if not isinstance(data, list):
                            logger.error(f"文件 {file_path} 的内容不是有效的 JSON 列表")
                            continue

                        for item in data:
                            question = item.get("question", "")
                            answer = item.get("answer", "")

                            # 验证数据格式
                            if not question or not answer:
                                logger.warning(f"无效的数据项: {item}")
                                continue

                            # 创建 Document 对象，将 question 作为主要内容，answer 和其他信息作为元数据
                            doc = Document(
                                page_content=question,
                                metadata={
                                    "source": file_path,
                                    "answer": answer,
                                    "type": "faq"
                                }
                            )
                            documents.append(doc)
                except Exception as e:
                    logger.error(f"加载文件 {file_path} 失败：{e}")
                    continue

        if documents:
            # 自定义文档解析器
            if self.docs_resolver:
                documents = self.docs_resolver(documents)

            # 添加到向量数据库
            self.retriever.add_documents(documents)
            logger.info(f"成功添加 {len(documents)} 个文档到向量数据库")
        else:
            logger.warning("没有找到任何文档")

    def split_documents(self, file_path) -> List[Document]:
        """
        读取文件并转换为document对象，该方法可以被子类重写
        """
        documents = []
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            # documents.append(Document(page_content=content, metadata={"source": file_path}))
            # 分割文档
            split_docs = self.text_splitter.split_text(content)
            documents.extend([Document(page_content=doc, metadata={"source": file_path}) for doc in split_docs])
        if not split_docs:
            logger.warning("分割后的文档为空")
            return documents
        return documents


def docs_resolve(documents: List[Document]) -> List[Document]:
    """
    解析文档列表，返回一个新的文档列表。
    这里可以添加更多的解析逻辑，在不同的子类中可以重写此方法以实现特定的解析逻辑。
    """
    if not documents:
        return []
    return documents


class CustomTextSplitter(CharacterTextSplitter):
    """
    自定义文本分割器，继承自CharacterTextSplitter。
    主要用于将文本分割成更小的块，以便于后续处理。
    """

    def __init__(self, chunk_size: int = 100, chunk_overlap: int = 0):
        super().__init__(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    def split_text(self, text: str) -> list:
        """
        重写文本分割方法，使用简单的字符分割。
        """
        lines = [line for line in text.strip().split("\n") if line.strip()]
        return lines if lines else []


def doc_resolve(documents: List[Document]) -> List[Document]:
    """
    解析文档列表，返回一个新的文档列表，将数据按|进行拆分
    """
    if not documents:
        return []
    resolved_documents = []
    for doc in documents:
        # 将文档内容按'|'分割,例如 "question|answer" ,answer部分我们放到metadata中
        lines = doc.page_content.strip().split("\n")
        for line in [x for x in lines if x.strip()]:
            parts = line.split('|')
            if len(parts) < 2:
                logger.warning(f"文档内容格式不正确: {line}")
                continue
            doc.metadata["answer"] = parts[1].strip()
            resolved_documents.append(Document(page_content=parts[0].strip(), metadata=doc.metadata))
    return resolved_documents
